import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader, Dataset
import tkinter as tk
from tkinter import simpledialog, messagebox
from PIL import Image, ImageDraw
import os
import numpy as np
import pickle

# 确保数据目录存在
if not os.path.exists('./data'):
    os.makedirs('./data')
USER_DATA_FILE = './data/user_handwritten_data.pkl'
MODEL_FILE = './data/model.pth'  # 模型保存路径

# 定义数据预处理
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])


# 自定义数据集用于存储用户手写数据
class UserDataset(Dataset):
    def __init__(self, data=None, transform=None):
        self.data = data if data is not None else []
        self.transform = transform

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        image, label = self.data[idx]
        if self.transform:
            image = self.transform(image)
        return image, label

    def add_sample(self, image, label):
        image_28 = image.resize((28, 28), Image.LANCZOS)
        self.data.append((image_28, label))

    def save_to_file(self, filename):
        with open(filename, 'wb') as f:
            pickle.dump(self.data, f)

    def load_from_file(self, filename):
        if os.path.exists(filename):
            with open(filename, 'rb') as f:
                self.data = pickle.load(f)


# 加载MNIST训练集和测试集
trainset = torchvision.datasets.MNIST(root='./data', train=True,
                                      download=True, transform=transform)
testset = torchvision.datasets.MNIST(root='./data', train=False,
                                     download=True, transform=transform)

# 加载用户数据
user_dataset = UserDataset(transform=transform)
user_dataset.load_from_file(USER_DATA_FILE)

# 合并原始训练集和用户数据集
combined_dataset = torch.utils.data.ConcatDataset([trainset, user_dataset])

# 创建数据加载器
trainloader = DataLoader(combined_dataset, batch_size=64, shuffle=True)
testloader = DataLoader(testset, batch_size=64, shuffle=False)


# 定义卷积神经网络模型
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.fc1 = nn.Linear(20 * 4 * 4, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        x = self.pool(torch.relu(self.conv1(x)))
        x = self.pool(torch.relu(self.conv2(x)))
        x = x.view(-1, 20 * 4 * 4)
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x


# 训练模型
def train_model(model, trainloader, criterion, optimizer, epochs):
    model.train()
    for epoch in range(epochs):
        running_loss = 0.0
        for i, data in enumerate(trainloader, 0):
            inputs, labels = data
            optimizer.zero_grad()

            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            running_loss += loss.item()
            if i % 200 == 199:
                print(f'[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 200:.3f}')
                running_loss = 0.0


# 测试模型
def test_model(model, testloader):
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for data in testloader:
            images, labels = data
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    accuracy = 100 * correct / total
    print(f'Accuracy of the network on the 10000 test images: {accuracy}%')
    return accuracy


# 加载或训练模型
def load_or_train_model():
    model = SimpleCNN()
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.5)

    if os.path.exists(MODEL_FILE):
        print("加载已保存的模型...")
        model.load_state_dict(torch.load(MODEL_FILE))
        # 测试加载的模型
        accuracy = test_model(model, testloader)
    else:
        print("初始模型训练中...")
        train_model(model, trainloader, criterion, optimizer, epochs=5)
        accuracy = test_model(model, testloader)
        # 保存训练好的模型
        torch.save(model.state_dict(), MODEL_FILE)
        print("模型已保存！")
    return model, criterion, optimizer, accuracy


# 创建一个Tkinter窗口
root = tk.Tk()
root.title("手写数字识别")

# 加载或训练模型
model, criterion, optimizer, initial_accuracy = load_or_train_model()

# 显示模型准确率
accuracy_label = tk.Label(root, text=f"模型准确率: {initial_accuracy:.2f}%")
accuracy_label.pack()

# 创建一个28x28的画布（实际显示为280x280以方便手写）
canvas_width = 280
canvas_height = 280
canvas = tk.Canvas(root, width=canvas_width, height=canvas_height, bg="white")
canvas.pack()

# 绘制表格
for i in range(1, 28):
    canvas.create_line(i * 10, 0, i * 10, canvas_height, fill="gray")
    canvas.create_line(0, i * 10, canvas_width, i * 10, fill="gray")

# 创建一个PIL图像用于保存手写数据
image = Image.new("L", (280, 280), 0)
draw = ImageDraw.Draw(image)


# 鼠标拖动事件处理函数
def paint(event):
    x1, y1 = (event.x - 5), (event.y - 5)
    x2, y2 = (event.x + 5), (event.y + 5)
    canvas.create_oval(x1, y1, x2, y2, fill="black", width=10)
    draw.ellipse([x1, y1, x2, y2], fill=255, width=10)


canvas.bind("<B1-Motion>", paint)


# 询问用户正确标签并添加到训练集
def correct_and_add(recognized_label, user_image):
    correct_label = simpledialog.askinteger(
        "纠正识别结果",
        f"系统识别为: {recognized_label}\n请输入正确的数字 (0-9):",
        minvalue=0, maxvalue=9
    )

    if correct_label is not None:
        user_dataset.add_sample(user_image, correct_label)
        user_dataset.save_to_file(USER_DATA_FILE)

        global combined_dataset, trainloader
        combined_dataset = torch.utils.data.ConcatDataset([trainset, user_dataset])
        trainloader = DataLoader(combined_dataset, batch_size=64, shuffle=True)

        messagebox.showinfo("成功", f"已将数字 {correct_label} 添加到训练集！")
        return True
    return False


# 识别按钮点击事件处理函数
def recognize():
    resized_image = image.resize((28, 28), Image.LANCZOS)
    input_tensor = transform(resized_image).unsqueeze(0)
    model.eval()
    with torch.no_grad():
        output = model(input_tensor)
        _, predicted = torch.max(output.data, 1)
        confidence = torch.nn.functional.softmax(output, dim=1)[0][predicted.item()].item() * 100

    predicted_label = predicted.item()
    result_label.config(text=f"预测结果: {predicted_label} (置信度: {confidence:.2f}%)")

    if confidence < 70:
        correct = messagebox.askyesno("确认", f"系统识别为 {predicted_label}，是否正确？")
        if not correct:
            correct_and_add(predicted_label, image)
    else:
        if messagebox.askyesno("确认", f"系统识别为 {predicted_label}，是否正确？"):
            if messagebox.askyesno("添加训练数据", "是否将此样本添加到训练集以改进模型？"):
                user_dataset.add_sample(image, predicted_label)
                user_dataset.save_to_file(USER_DATA_FILE)
                messagebox.showinfo("成功", f"已将数字 {predicted_label} 添加到训练集！")
        else:
            correct_and_add(predicted_label, image)


# 重新训练模型
def retrain_model():
    if len(user_dataset) == 0:
        messagebox.showinfo("提示", "没有新的用户数据需要训练！")
        return

    messagebox.showinfo("开始训练", f"使用 {len(user_dataset)} 个用户样本进行重新训练...")
    train_model(model, trainloader, criterion, optimizer, epochs=2)
    new_accuracy = test_model(model, testloader)
    accuracy_label.config(text=f"模型准确率: {new_accuracy:.2f}%")
    # 保存重新训练后的模型
    torch.save(model.state_dict(), MODEL_FILE)
    messagebox.showinfo("训练完成", f"重新训练完成！新的准确率: {new_accuracy:.2f}%，模型已保存。")


# 清除按钮点击事件处理函数
def clear():
    canvas.delete("all")
    draw.rectangle([0, 0, 280, 280], fill=0)
    result_label.config(text="预测结果: ")
    for i in range(1, 28):
        canvas.create_line(i * 10, 0, i * 10, canvas_height, fill="gray")
        canvas.create_line(0, i * 10, canvas_width, i * 10, fill="gray")


# 创建按钮框架
button_frame = tk.Frame(root)
button_frame.pack(pady=5)

# 创建识别按钮
recognize_button = tk.Button(button_frame, text="识别", command=recognize)
recognize_button.pack(side=tk.LEFT, padx=5)

# 创建重新训练按钮
retrain_button = tk.Button(button_frame, text="重新训练模型", command=retrain_model)
retrain_button.pack(side=tk.LEFT, padx=5)

# 创建清除按钮
clear_button = tk.Button(button_frame, text="清除", command=clear)
clear_button.pack(side=tk.LEFT, padx=5)

# 创建显示预测结果的标签
result_label = tk.Label(root, text="预测结果: ")
result_label.pack(pady=5)

# 显示用户数据统计
data_count_label = tk.Label(root, text=f"已收集用户样本数: {len(user_dataset)}")
data_count_label.pack(pady=5)


# 更新数据计数标签
def update_data_count():
    data_count_label.config(text=f"已收集用户样本数: {len(user_dataset)}")
    root.after(1000, update_data_count)


update_data_count()

root.mainloop()